/****************************************************************************
 * Copyright (c) 2018-2023 by the Cabana authors                            *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Cabana library. Cabana is distributed under a   *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#include <Cabana_Grid.hpp>

#include <Kokkos_Core.hpp>

#include <iostream>

#include <mpi.h>

//---------------------------------------------------------------------------//
// Array example.
//---------------------------------------------------------------------------//
void arrayExample()
{
    /*
      An array layout includes the local grid size information together with
      the number of values per mesh entity in the field (dimensionality), and
      the entity type on which the field is defined. An array stores the field
      data itself and may only be defined on a single entity type.
    */

    using exec_space = Kokkos::DefaultHostExecutionSpace;
    using device_type = exec_space::device_type;

    /*
      The array layout is an intermediate class between the local grid and array
      - it does not store the mesh field data. Instead it holds the layout of
      the field arrays.

      As with the previous examples, we define everything necessary to create
      the local grid.
    */

    // Let MPI compute the partitioning for this test.
    Cabana::Grid::DimBlockPartitioner<3> partitioner;

    // Create the global mesh.
    double cell_size = 0.50;
    std::array<int, 3> global_num_cell = { 37, 15, 20 };
    std::array<bool, 3> is_dim_periodic = { true, true, true };
    std::array<double, 3> global_low_corner = { 1.2, 3.3, -2.8 };
    std::array<double, 3> global_high_corner = {
        global_low_corner[0] + cell_size * global_num_cell[0],
        global_low_corner[1] + cell_size * global_num_cell[1],
        global_low_corner[2] + cell_size * global_num_cell[2] };
    auto global_mesh = Cabana::Grid::createUniformGlobalMesh(
        global_low_corner, global_high_corner, global_num_cell );

    // Create the global grid.
    auto global_grid = Cabana::Grid::createGlobalGrid(
        MPI_COMM_WORLD, global_mesh, is_dim_periodic, partitioner );

    // Get the current rank for printing output.
    int comm_rank = global_grid->blockId();
    if ( comm_rank == 0 )
    {
        std::cout << "Cabana::Grid Array Example" << std::endl;
        std::cout << "    (intended to be run with MPI)\n" << std::endl;
    }

    /*
      An array layout includes the local grid size information together with the
      dimensionality of the field data, the degrees of freedom on the mesh.
    */
    int halo_width = 2;
    int dofs_per_node = 4;
    auto node_layout = Cabana::Grid::createArrayLayout(
        global_grid, halo_width, dofs_per_node, Cabana::Grid::Node() );
    int view_rank = node_layout->num_space_dim + 1;

    /*
      Similar to the local grid, the array layout holds index spaces for
      parallel iteration over mesh fields. These again include options for mesh
      entities and owned/owned+ghosted.

      Note that index spaces generated by array layouts have an additional
      dimension over the field values on each entity.
    */
    auto array_node_owned_space =
        node_layout->indexSpace( Cabana::Grid::Own(), Cabana::Grid::Local() );
    std::cout << "Array layout (Own, Local) \nMin: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_owned_space.min( d ) << " ";
    std::cout << "\nMax: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_owned_space.max( d ) << " ";
    std::cout << "\n" << std::endl;

    auto array_node_ghosted_space =
        node_layout->indexSpace( Cabana::Grid::Ghost(), Cabana::Grid::Local() );
    std::cout << "Array layout (Ghost, Local) \nMin: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_ghosted_space.min( d ) << " ";
    std::cout << "\nMax: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_ghosted_space.max( d ) << " ";
    std::cout << "\n" << std::endl;

    auto array_node_shared_owned_space =
        node_layout->sharedIndexSpace( Cabana::Grid::Own(), -1, 0, 1 );
    std::cout << "Array layout (Shared, Own, -1,0,1) \nMin: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_shared_owned_space.min( d ) << " ";
    std::cout << "\nMax: ";
    for ( int d = 0; d < view_rank; ++d )
        std::cout << array_node_shared_owned_space.max( d ) << " ";
    std::cout << "\n" << std::endl;

    /*
      Next, we create an array from the array layout. This directly stores the
      Kokkos View with field data that can be extracted for parallel
      computation.
    */
    std::string label( "example_array" );
    auto array =
        Cabana::Grid::createArray<double, device_type>( label, node_layout );

    auto view = array->view();
    std::cout << "Array total size: " << view.size() << std::endl;
    std::cout << "Array extents: ";
    for ( int i = 0; i < view_rank; ++i )
        std::cout << view.extent( i ) << " ";
    std::cout << "\n" << std::endl;

    /*
      Many array-based operations are available for convenience:
       - assign a value to every field on the mesh
       - scale the entire array
       - update
       - dot product between two arrays
       - two-norm
       - one-norm
       - infinity-norm

      For each operation we print the first value in the array - each value is
      updated.
    */
    Cabana::Grid::ArrayOp::assign( *array, 2.0, Cabana::Grid::Ghost() );

    // Scale the entire array with a single value.
    Cabana::Grid::ArrayOp::scale( *array, 0.5, Cabana::Grid::Ghost() );

    /*
      Compute the dot product of the two arrays.

      Note that in this case, two arrays are updated in a single operation. It
      is also possible to update three in one operation.
    */
    std::vector<double> dots( dofs_per_node );
    auto array_2 =
        Cabana::Grid::createArray<double, device_type>( label, node_layout );
    Cabana::Grid::ArrayOp::assign( *array_2, 0.5, Cabana::Grid::Ghost() );
    Cabana::Grid::ArrayOp::update( *array, 3.0, *array_2, 2.0,
                                   Cabana::Grid::Ghost() );
    Cabana::Grid::ArrayOp::dot( *array, *array_2, dots );
    std::cout << "Array dot product: ";
    std::cout << dots[0] << std::endl;

    // Compute the two-norm of the array components
    std::cout << "Array two-norm: ";
    std::vector<double> norm_2( dofs_per_node );
    Cabana::Grid::ArrayOp::norm2( *array, norm_2 );
    std::cout << norm_2[0] << std::endl;

    // Compute the one-norm of the array components
    std::cout << "Array one-norm: ";
    std::vector<double> norm_1( dofs_per_node );
    Cabana::Grid::ArrayOp::norm1( *array, norm_1 );
    std::cout << norm_1[0] << std::endl;

    // Compute the infinity-norm of the array components
    std::cout << "Array infinity-norm: ";
    view = array->view();
    std::vector<double> large_vals = { -1939304932.2, 20399994.532,
                                       9098201010.114, -89877402343.99 };
    for ( int n = 0; n < dofs_per_node; ++n )
        view( 4, 4, 4, n ) = large_vals[n];
    std::vector<double> norm_inf( dofs_per_node );
    Cabana::Grid::ArrayOp::normInf( *array, norm_inf );
    std::cout << norm_inf[0] << std::endl;

    /*
      It is also possible to copy arrays with compatible index spaces.
    */
    Cabana::Grid::ArrayOp::copy( *array, *array_2, Cabana::Grid::Own() );

    /*
      Cloning an array makes a new, identical array to the given array with a
      new memory allocation.
    */
    auto array_3 = Cabana::Grid::ArrayOp::clone( *array );
    Cabana::Grid::ArrayOp::copy( *array_3, *array, Cabana::Grid::Own() );

    /*
      A clone and copy can also be done in a single operation.
    */
    auto array_4 =
        Cabana::Grid::ArrayOp::cloneCopy( *array, Cabana::Grid::Own() );

    /*
      It is also possible to create sub-arrays of existing arrays.
    */
    auto subarray = Cabana::Grid::createSubarray( *array, 2, 4 );
    auto sub_ghosted_space = subarray->layout()->indexSpace(
        Cabana::Grid::Ghost(), Cabana::Grid::Local() );
    std::cout << "\nSub-array index space size: ";
    std::cout << sub_ghosted_space.size() << std::endl;
    std::cout << "Sub-array index space extents: ";
    for ( int n = 0; n < dofs_per_node; ++n )
        std::cout << sub_ghosted_space.extent( n ) << " ";
    std::cout << "\n" << std::endl;
}

//---------------------------------------------------------------------------//
// Main.
//---------------------------------------------------------------------------//
int main( int argc, char* argv[] )
{
    MPI_Init( &argc, &argv );
    {
        Kokkos::ScopeGuard scope_guard( argc, argv );

        arrayExample();
    }
    MPI_Finalize();

    return 0;
}

//---------------------------------------------------------------------------//
